/**
 ******************************************************************************
 * Copyright (c) 2022 - ~, SCUT-RobotLab Development Team
 * @file    RobotRunner.cpp
 * @author
 * @brief   Code for .
 * @date   2022-02-28
 * @version 1.0
 * @par Change Log:
 * <table>
 * <tr><th>Date        <th>Version  <th>Author     <th>Description
 * </table>
 *
 ******************************************************************************
 * @attention
 *
 * if you had modified this file, please make sure your code does not have many
 * bugs, update the version Number, write dowm your name and the date, the most
 * important is make sure the users will have clear and definite understanding
 * through your new brief.
 *
 * <h2><center>&copy; Copyright (c) 2022 - ~,SCUT-RobotLab Development Team.
 * All rights reserved.</center></h2>
 ******************************************************************************
 */
#include <robot/RobotRunner.h>
#include "common/Dynamics/Zouwu.h"
#include "common/Controllers/PositionVelocityEstimator.h"
#include "common/Controllers/ContactEstimator.h"
#include "common/Controllers/OrientationEstimator.h"

RobotRunner::RobotRunner(RobotController* robot_ctrl) {
    _robot_ctrl = robot_ctrl;

    // Initial params
    controlParameters.initializeFromYamlFile("robot_parameters");
}

/**
 * Initializes the robot model, state estimator, leg controller,
 * robot data, and any control logic specific data.
 */
void RobotRunner::init() {
    printf("[RobotRunner] initialize\n");

    // Build the appropriate Quadruped object
    _quadruped = buildZouwu<float>();

    // Initialize robot data
    _model = _quadruped.buildModel();
    _jpos_initializer = new JPosInitializer<float>(3., controlParameters.controller_dt);

    // initialize the leg controller
    _legController = new LegController<float>(_quadruped);
    _stateEstimator = new StateEstimatorContainer<float>(&cheaterState, &vectorNavData, _legController->datas,
                                                         &_stateEstimate, &controlParameters);
    initializeStateEstimator();

    // test
    _stateEstimator_test = new StateEstimatorContainer<float>(&cheaterState, &vectorNavData, _legController->datas,
                                                         &_stateEstimate_test, &controlParameters);
    _stateEstimator_test->removeAllEstimators();
    _stateEstimator_test->addEstimator<ContactEstimator<float>>();
    Vec4<float> contactDefault;
    contactDefault << 0.5, 0.5, 0.5, 0.5;
    _stateEstimator_test->setContactPhase(contactDefault);
    _stateEstimator_test->addEstimator<VectorNavOrientationEstimator<float>>();
    _stateEstimator_test->addEstimator<LinearKFPositionVelocityEstimator<float>>();


    // Initialize the DesiredStateCommand object
    _desiredStateCommand = new DesiredStateCommand<float>(&rc_control, &controlParameters, &_stateEstimate,
                                                          controlParameters.controller_dt);

    // Controller initializations
    _robot_ctrl->_model = &_model;
    _robot_ctrl->_quadruped = &_quadruped;
    _robot_ctrl->_stateEstimator = _stateEstimator;
    _robot_ctrl->_stateEstimate = &_stateEstimate;
    _robot_ctrl->_legController = _legController;
    _robot_ctrl->_controlParameters = &controlParameters;
    _robot_ctrl->_desiredStateCommand = _desiredStateCommand;

    _robot_ctrl->initializeController();
}

/**
 * Runs the overall robot control system by calling each of the major components
 * to run each of their respective steps.
 */
RobotCommand* RobotRunner::run() {
    // Run the state estimator step
    if(controlParameters.cheater_mode)
    _stateEstimator->run();

    // Update the data from the robot
    setupStep();

    static int count_ini(0);
    ++count_ini;
    if ((rc_control.mode == RC_mode::OFF) && controlParameters.use_rc) {
        if (count_ini % 1000 == 0)
            printf("ESTOP!\n");
        _legController->zeroCommand();
        _robot_ctrl->Estop();
    } else {
        // Controller
        if (!_jpos_initializer->IsInitialized(_legController)) {
            Mat3<float> kpMat;
            Mat3<float> kdMat;
            // Update the jpos feedback gains
            kpMat << 50, 0, 0, 0, 50, 0, 0, 0, 50;
            kdMat << 1, 0, 0, 0, 1, 0, 0, 0, 1;
            for (int leg = 0; leg < 4; leg++) {
                _legController->commands[leg].kpJoint = kpMat;
                _legController->commands[leg].kdJoint = kdMat;
            }
        } else {
            // Run Control
            _robot_ctrl->runController();
        }
    }
 
    // Sets the leg controller commands for the robot appropriate commands
    finalizeStep();

    // test 
    // Vec3<float> p(0,0,-0.25);
    // Vec3<float> q;
    // computeAuxiliaryAngle(_quadruped, p, &q);
    // for(int leg = 0; leg < 4; leg ++){
    //     robot_cmd.q_des_abad[leg]= q[0];
    //     robot_cmd.q_des_hip[leg] = q[1] + q[2] ;
    //     robot_cmd.q_des_knee[leg] = q[2] - q[1];

    //     robot_cmd.mode_abad[leg] = 0;
    //     robot_cmd.mode_hip[leg] = 0;
    //     robot_cmd.mode_knee[leg] = 0;     
    // }

    // control command out
    if (_legController->_legsEnabled)
        return &robot_cmd;
    else
        return nullptr;
}

/*!
 * Before running user code, setup the leg control and estimators
 */
void RobotRunner::setupStep() {
    // Setup the leg controller for a new iteration
    _legController->zeroCommand();
    // _legController->setEnabled(true);

  if (!_cheaterModeEnabled && controlParameters.cheater_mode) {
    printf("[RobotRunner] Transitioning to Cheater Mode...\n");
    initializeStateEstimator(true);
    // todo any configuration
    _cheaterModeEnabled = true;
  }

  // check transition from cheater mode:
  if (_cheaterModeEnabled && !controlParameters.cheater_mode) {
    printf("[RobotRunner] Transitioning from Cheater Mode...\n");
    initializeStateEstimator(false);
    // todo any configuration
    _cheaterModeEnabled = false;
  }

    // todo safety checks, sanity checks, etc...
}

/*!
 * After the user code, send leg commands, update state estimate, and publish debug data
 */
void RobotRunner::finalizeStep() {
    // check data vailty
    uint8_t invaild_num = 0;
    if(rc_control.mode == RC_mode::LOCOMOTION && ctrl_mode == Mode_Tau){
        for(int leg = 0; leg < 4; leg++){
            if(robot_cmd.mode_hip[0]== Mode_Tau&& _stateEstimator->getResult().contactEstimate(leg) > 0 
                    && fabs(_legController->commands[leg].forceFeedForward[2]) < 0.1f){
                invaild_num ++;
            }
        }
        if(invaild_num >= 2){
            robot_cmd = last_cmd;
        }
        else{
            _legController->updateCommand(&robot_cmd);
            last_cmd = robot_cmd;
        }
    }
    else _legController->updateCommand(&robot_cmd); 
    _iterations++;
}

/*!
 * Reset the state estimator in the given mode.
 * @param cheaterMode
 */
void RobotRunner::initializeStateEstimator(bool cheaterMode) {
    _stateEstimator->removeAllEstimators();
    _stateEstimator->addEstimator<ContactEstimator<float>>();
    Vec4<float> contactDefault;
    contactDefault << 0.5, 0.5, 0.5, 0.5;
    _stateEstimator->setContactPhase(contactDefault);
    if (cheaterMode) {
        _stateEstimator->addEstimator<CheaterOrientationEstimator<float>>();
        _stateEstimator->addEstimator<CheaterPositionVelocityEstimator<float>>();
    } else {
        _stateEstimator->addEstimator<VectorNavOrientationEstimator<float>>();
        _stateEstimator->addEstimator<LinearKFPositionVelocityEstimator<float>>();
    }
}

void RobotRunner::updateRC(const rc_control_settings* data) { rc_control = *data; }

void RobotRunner::updateRCvel(float vel_x, float vel_y, float omega){
    rc_control.v_des[0] = fminf(fmaxf(vel_x, -0.3), 0.3) ;
    rc_control.v_des[1] = fminf(fmaxf(vel_y, -0.3), 0.3) ;
    rc_control.omega_des[2] = fminf(fmaxf(omega, -0.3), 0.3) ;
}

void RobotRunner::updateRCgait(uint8_t gait){
    rc_control.variable[0] = gait;
}

void RobotRunner::updataCheaterState(const CheaterState<float>* state){
    // cheaterState.acceleration = state->acceleration;
    // cheaterState.omegaBody = state->omegaBody;
    // cheaterState.orientation = state->orientation;
    // cheaterState.position = state->position;
    // cheaterState.vBody = state->vBody;
    cheaterState = *state;
}

void RobotRunner::updateVectorNavData(const VectorNavData* nav){ vectorNavData = *nav;}

RobotRunner::~RobotRunner() {
    delete _legController;
    delete _stateEstimator;
    delete _jpos_initializer;
}

void RobotRunner::cleanup() {}

/************************ COPYRIGHT(C) SCUT-RobotLab Development Team **************************/
